#!/usr/bin/env python
# coding=UTF-8
"""
@Author: Tao Hang
@LastEditors: Tao Hang
@Description:
@Date: 2019-03-26 10:30:19
@LastEditTime: 2019-04-15 11:25:16
"""

import numpy as np
import torch
import copy
from torch.autograd import Variable

from .attack import Attack


class BIM(Attack):
    def __init__(self, model=None, device=None, IsTargeted=None, config=None):
        """
        @description: Fast Gradient Sign Method (FGSM)
        @param {
            model:需要测试的模型
            device: 设备(GPU)
            IsTargeted:是否是目标攻击
            kwargs: 用户对攻击方法需要的参数
        }
        @return: None
        """
        super(BIM, self).__init__(model, device, IsTargeted)
        # 使用该函数时候，要保证训练模型的标签是从0开始，而不是1
        self.criterion = torch.nn.CrossEntropyLoss()
        self._parse_params(config)

    def _parse_params(self, config):
        """
        @description:
        @param {
            epsilon:沿着梯度方向步长的参数
        }
        @return: None
        """
        self.eps = float(config.get("epsilon", 0.03))
        self.eps_iter = float(config.get("eps_iter", 0.002))
        self.num_steps = int(config.get("num_steps", 15))

    def generate(self, images, labels):
        """
        @description:
        @param {
            xs:原始的样本
            ys:样本的标签
        }
        @return: adv_xs{numpy.ndarray}
        """
        device = self.device
        targeted = self.IsTargeted

        var_images = copy.deepcopy(images)
        var_images = var_images.clone().detach()
        labels = labels.to(device)

        for _loopnum in range(self.num_steps):
            var_images = var_images.to(device)
            var_images.requires_grad_(True)
            outputs = self.model(var_images)
            if targeted:
                loss = -self.criterion(outputs, labels)
            else:
                loss = self.criterion(outputs, labels)

            grad = torch.autograd.grad(loss, var_images, retain_graph=False, create_graph=False)[0]

            adv_noise = self.eps_iter * grad.sign()
            var_images = var_images + torch.clamp(adv_noise, min=-self.eps, max=self.eps)
            var_images = torch.clamp(var_images, min=0, max=1)
            var_images.detach_()

        return var_images
